# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #

from math import nan

from gpu.host import DeviceContext
from memory import bitcast


# CHECK: 0.0, 7.62939453125e-06, 1.52587890625e-05, 2.288818359375e-05, 3.0517578125e-05, 3.814697265625e-05, 4.57763671875e-05, 5.340576171875e-05,
# CHECK: 6.103515625e-05, 7.62939453125e-05, 9.1552734375e-05, 0.0001068115234375, 0.0001220703125, 0.000152587890625, 0.00018310546875, 0.000213623046875,
# CHECK: 0.000244140625, 0.00030517578125, 0.0003662109375, 0.00042724609375, 0.00048828125, 0.0006103515625, 0.000732421875, 0.0008544921875,
# CHECK: 0.0009765625, 0.001220703125, 0.00146484375, 0.001708984375, 0.001953125, 0.00244140625, 0.0029296875, 0.00341796875,
# CHECK: 0.00390625, 0.0048828125, 0.005859375, 0.0068359375, 0.0078125, 0.009765625, 0.01171875, 0.013671875,
# CHECK: 0.015625, 0.01953125, 0.0234375, 0.02734375, 0.03125, 0.0390625, 0.046875, 0.0546875,
# CHECK: 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.15625, 0.1875, 0.21875,
# CHECK: 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.625, 0.75, 0.875,
# CHECK: 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 3.5,
# CHECK: 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0,
# CHECK: 16.0, 20.0, 24.0, 28.0, 32.0, 40.0, 48.0, 56.0,
# CHECK: 64.0, 80.0, 96.0, 112.0, 128.0, 160.0, 192.0, 224.0,
# CHECK: 256.0, 320.0, 384.0, 448.0, 512.0, 640.0, 768.0, 896.0,
# CHECK: 1024.0, 1280.0, 1536.0, 1792.0, 2048.0, 2560.0, 3072.0, 3584.0,
# CHECK: 4096.0, 5120.0, 6144.0, 7168.0, 8192.0, 10240.0, 12288.0, 14336.0,
# CHECK: 16384.0, 20480.0, 24576.0, 28672.0, 32768.0, 40960.0, 49152.0, 57344.0,
# CHECK: nan, -7.62939453125e-06, -1.52587890625e-05, -2.288818359375e-05, -3.0517578125e-05, -3.814697265625e-05, -4.57763671875e-05, -5.340576171875e-05,
# CHECK: -6.103515625e-05, -7.62939453125e-05, -9.1552734375e-05, -0.0001068115234375, -0.0001220703125, -0.000152587890625, -0.00018310546875, -0.000213623046875,
# CHECK: -0.000244140625, -0.00030517578125, -0.0003662109375, -0.00042724609375, -0.00048828125, -0.0006103515625, -0.000732421875, -0.0008544921875,
# CHECK: -0.0009765625, -0.001220703125, -0.00146484375, -0.001708984375, -0.001953125, -0.00244140625, -0.0029296875, -0.00341796875,
# CHECK: -0.00390625, -0.0048828125, -0.005859375, -0.0068359375, -0.0078125, -0.009765625, -0.01171875, -0.013671875,
# CHECK: -0.015625, -0.01953125, -0.0234375, -0.02734375, -0.03125, -0.0390625, -0.046875, -0.0546875,
# CHECK: -0.0625, -0.078125, -0.09375, -0.109375, -0.125, -0.15625, -0.1875, -0.21875,
# CHECK: -0.25, -0.3125, -0.375, -0.4375, -0.5, -0.625, -0.75, -0.875,
# CHECK: -1.0, -1.25, -1.5, -1.75, -2.0, -2.5, -3.0, -3.5,
# CHECK: -4.0, -5.0, -6.0, -7.0, -8.0, -10.0, -12.0, -14.0,
# CHECK: -16.0, -20.0, -24.0, -28.0, -32.0, -40.0, -48.0, -56.0,
# CHECK: -64.0, -80.0, -96.0, -112.0, -128.0, -160.0, -192.0, -224.0,
# CHECK: -256.0, -320.0, -384.0, -448.0, -512.0, -640.0, -768.0, -896.0,
# CHECK: -1024.0, -1280.0, -1536.0, -1792.0, -2048.0, -2560.0, -3072.0, -3584.0,
# CHECK: -4096.0, -5120.0, -6144.0, -7168.0, -8192.0, -10240.0, -12288.0, -14336.0,
# CHECK: -16384.0, -20480.0, -24576.0, -28672.0, -32768.0, -40960.0, -49152.0, -57344.0,
fn test_e5m2fnuz_initialization():
    print("== test_e5m2fnuz_initialization")

    var simd_e5m2fnuz = SIMD[DType.float8_e5m2fnuz, 256](
        0.0,
        7.62939453125e-06,
        1.52587890625e-05,
        2.288818359375e-05,
        3.0517578125e-05,
        3.814697265625e-05,
        4.57763671875e-05,
        5.340576171875e-05,
        6.103515625e-05,
        7.62939453125e-05,
        9.1552734375e-05,
        0.0001068115234375,
        0.0001220703125,
        0.000152587890625,
        0.00018310546875,
        0.000213623046875,
        0.000244140625,
        0.00030517578125,
        0.0003662109375,
        0.00042724609375,
        0.00048828125,
        0.0006103515625,
        0.000732421875,
        0.0008544921875,
        0.0009765625,
        0.001220703125,
        0.00146484375,
        0.001708984375,
        0.001953125,
        0.00244140625,
        0.0029296875,
        0.00341796875,
        0.00390625,
        0.0048828125,
        0.005859375,
        0.0068359375,
        0.0078125,
        0.009765625,
        0.01171875,
        0.013671875,
        0.015625,
        0.01953125,
        0.0234375,
        0.02734375,
        0.03125,
        0.0390625,
        0.046875,
        0.0546875,
        0.0625,
        0.078125,
        0.09375,
        0.109375,
        0.125,
        0.15625,
        0.1875,
        0.21875,
        0.25,
        0.3125,
        0.375,
        0.4375,
        0.5,
        0.625,
        0.75,
        0.875,
        1.0,
        1.25,
        1.5,
        1.75,
        2.0,
        2.5,
        3.0,
        3.5,
        4.0,
        5.0,
        6.0,
        7.0,
        8.0,
        10.0,
        12.0,
        14.0,
        16.0,
        20.0,
        24.0,
        28.0,
        32.0,
        40.0,
        48.0,
        56.0,
        64.0,
        80.0,
        96.0,
        112.0,
        128.0,
        160.0,
        192.0,
        224.0,
        256.0,
        320.0,
        384.0,
        448.0,
        512.0,
        640.0,
        768.0,
        896.0,
        1024.0,
        1280.0,
        1536.0,
        1792.0,
        2048.0,
        2560.0,
        3072.0,
        3584.0,
        4096.0,
        5120.0,
        6144.0,
        7168.0,
        8192.0,
        10240.0,
        12288.0,
        14336.0,
        16384.0,
        20480.0,
        24576.0,
        28672.0,
        32768.0,
        40960.0,
        49152.0,
        57344.0,
        nan[DType.float8_e5m2fnuz](),
        -7.62939453125e-06,
        -1.52587890625e-05,
        -2.288818359375e-05,
        -3.0517578125e-05,
        -3.814697265625e-05,
        -4.57763671875e-05,
        -5.340576171875e-05,
        -6.103515625e-05,
        -7.62939453125e-05,
        -9.1552734375e-05,
        -0.0001068115234375,
        -0.0001220703125,
        -0.000152587890625,
        -0.00018310546875,
        -0.000213623046875,
        -0.000244140625,
        -0.00030517578125,
        -0.0003662109375,
        -0.00042724609375,
        -0.00048828125,
        -0.0006103515625,
        -0.000732421875,
        -0.0008544921875,
        -0.0009765625,
        -0.001220703125,
        -0.00146484375,
        -0.001708984375,
        -0.001953125,
        -0.00244140625,
        -0.0029296875,
        -0.00341796875,
        -0.00390625,
        -0.0048828125,
        -0.005859375,
        -0.0068359375,
        -0.0078125,
        -0.009765625,
        -0.01171875,
        -0.013671875,
        -0.015625,
        -0.01953125,
        -0.0234375,
        -0.02734375,
        -0.03125,
        -0.0390625,
        -0.046875,
        -0.0546875,
        -0.0625,
        -0.078125,
        -0.09375,
        -0.109375,
        -0.125,
        -0.15625,
        -0.1875,
        -0.21875,
        -0.25,
        -0.3125,
        -0.375,
        -0.4375,
        -0.5,
        -0.625,
        -0.75,
        -0.875,
        -1.0,
        -1.25,
        -1.5,
        -1.75,
        -2.0,
        -2.5,
        -3.0,
        -3.5,
        -4.0,
        -5.0,
        -6.0,
        -7.0,
        -8.0,
        -10.0,
        -12.0,
        -14.0,
        -16.0,
        -20.0,
        -24.0,
        -28.0,
        -32.0,
        -40.0,
        -48.0,
        -56.0,
        -64.0,
        -80.0,
        -96.0,
        -112.0,
        -128.0,
        -160.0,
        -192.0,
        -224.0,
        -256.0,
        -320.0,
        -384.0,
        -448.0,
        -512.0,
        -640.0,
        -768.0,
        -896.0,
        -1024.0,
        -1280.0,
        -1536.0,
        -1792.0,
        -2048.0,
        -2560.0,
        -3072.0,
        -3584.0,
        -4096.0,
        -5120.0,
        -6144.0,
        -7168.0,
        -8192.0,
        -10240.0,
        -12288.0,
        -14336.0,
        -16384.0,
        -20480.0,
        -24576.0,
        -28672.0,
        -32768.0,
        -40960.0,
        -49152.0,
        -57344.0,
    )

    for i in range(256):
        print(simd_e5m2fnuz[i], end=", ")
        if (i + 1) % 8 == 0:
            print("")


fn test_simd_e5m2fnuz_to_float[target: DType]():
    var float8_simd = SIMD[DType.float8_e5m2fnuz, 256](0.0)
    for i in range(256):
        float8_simd[i] = bitcast[DType.float8_e5m2fnuz](UInt8(i))

    target_casted = float8_simd.cast[target]()

    comptime M = 32
    comptime N = 8
    for i in range(M):
        for j in range(N):
            print(target_casted[i * N + j], end=", ")
        print("")


# CHECK-LABEL: test_simd_e5m2fnuz_to_f32
# CHECK: 0.0, 7.6293945e-06, 1.5258789e-05, 2.2888184e-05, 3.0517578e-05, 3.8146973e-05, 4.5776367e-05, 5.340576e-05,
# CHECK: 6.1035156e-05, 7.6293945e-05, 9.1552734e-05, 0.00010681152, 0.00012207031, 0.00015258789, 0.00018310547, 0.00021362305,
# CHECK: 0.00024414063, 0.00030517578, 0.00036621094, 0.0004272461, 0.00048828125, 0.00061035156, 0.0007324219, 0.0008544922,
# CHECK: 0.0009765625, 0.0012207031, 0.0014648438, 0.0017089844, 0.001953125, 0.0024414062, 0.0029296875, 0.0034179688,
# CHECK: 0.00390625, 0.0048828125, 0.005859375, 0.0068359375, 0.0078125, 0.009765625, 0.01171875, 0.013671875,
# CHECK: 0.015625, 0.01953125, 0.0234375, 0.02734375, 0.03125, 0.0390625, 0.046875, 0.0546875,
# CHECK: 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.15625, 0.1875, 0.21875,
# CHECK: 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.625, 0.75, 0.875,
# CHECK: 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 3.5,
# CHECK: 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0,
# CHECK: 16.0, 20.0, 24.0, 28.0, 32.0, 40.0, 48.0, 56.0,
# CHECK: 64.0, 80.0, 96.0, 112.0, 128.0, 160.0, 192.0, 224.0,
# CHECK: 256.0, 320.0, 384.0, 448.0, 512.0, 640.0, 768.0, 896.0,
# CHECK: 1024.0, 1280.0, 1536.0, 1792.0, 2048.0, 2560.0, 3072.0, 3584.0,
# CHECK: 4096.0, 5120.0, 6144.0, 7168.0, 8192.0, 10240.0, 12288.0, 14336.0,
# CHECK: 16384.0, 20480.0, 24576.0, 28672.0, 32768.0, 40960.0, 49152.0, 57344.0,
# CHECK: nan, -7.6293945e-06, -1.5258789e-05, -2.2888184e-05, -3.0517578e-05, -3.8146973e-05, -4.5776367e-05, -5.340576e-05,
# CHECK: -6.1035156e-05, -7.6293945e-05, -9.1552734e-05, -0.00010681152, -0.00012207031, -0.00015258789, -0.00018310547, -0.00021362305,
# CHECK: -0.00024414063, -0.00030517578, -0.00036621094, -0.0004272461, -0.00048828125, -0.00061035156, -0.0007324219, -0.0008544922,
# CHECK: -0.0009765625, -0.0012207031, -0.0014648438, -0.0017089844, -0.001953125, -0.0024414062, -0.0029296875, -0.0034179688,
# CHECK: -0.00390625, -0.0048828125, -0.005859375, -0.0068359375, -0.0078125, -0.009765625, -0.01171875, -0.013671875,
# CHECK: -0.015625, -0.01953125, -0.0234375, -0.02734375, -0.03125, -0.0390625, -0.046875, -0.0546875,
# CHECK: -0.0625, -0.078125, -0.09375, -0.109375, -0.125, -0.15625, -0.1875, -0.21875,
# CHECK: -0.25, -0.3125, -0.375, -0.4375, -0.5, -0.625, -0.75, -0.875,
# CHECK: -1.0, -1.25, -1.5, -1.75, -2.0, -2.5, -3.0, -3.5,
# CHECK: -4.0, -5.0, -6.0, -7.0, -8.0, -10.0, -12.0, -14.0,
# CHECK: -16.0, -20.0, -24.0, -28.0, -32.0, -40.0, -48.0, -56.0,
# CHECK: -64.0, -80.0, -96.0, -112.0, -128.0, -160.0, -192.0, -224.0,
# CHECK: -256.0, -320.0, -384.0, -448.0, -512.0, -640.0, -768.0, -896.0,
# CHECK: -1024.0, -1280.0, -1536.0, -1792.0, -2048.0, -2560.0, -3072.0, -3584.0,
# CHECK: -4096.0, -5120.0, -6144.0, -7168.0, -8192.0, -10240.0, -12288.0, -14336.0,
# CHECK: -16384.0, -20480.0, -24576.0, -28672.0, -32768.0, -40960.0, -49152.0, -57344.0,
fn test_simd_e5m2fnuz_to_f32():
    print("== test_simd_e5m2fnuz_to_f32")
    test_simd_e5m2fnuz_to_float[DType.float32]()


# CHECK-LABEL: test_simd_e5m2fnuz_to_f16
# CHECK: 0.0, 7.6293945e-06, 1.5258789e-05, 2.2888184e-05, 3.0517578e-05, 3.8146973e-05, 4.5776367e-05, 5.340576e-05,
# CHECK: 6.1035156e-05, 7.6293945e-05, 9.1552734e-05, 0.00010681152, 0.00012207031, 0.00015258789, 0.00018310547, 0.00021362305,
# CHECK: 0.00024414063, 0.00030517578, 0.00036621094, 0.0004272461, 0.00048828125, 0.00061035156, 0.0007324219, 0.0008544922,
# CHECK: 0.0009765625, 0.0012207031, 0.0014648438, 0.0017089844, 0.001953125, 0.0024414062, 0.0029296875, 0.0034179688,
# CHECK: 0.00390625, 0.0048828125, 0.005859375, 0.0068359375, 0.0078125, 0.009765625, 0.01171875, 0.013671875,
# CHECK: 0.015625, 0.01953125, 0.0234375, 0.02734375, 0.03125, 0.0390625, 0.046875, 0.0546875,
# CHECK: 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.15625, 0.1875, 0.21875,
# CHECK: 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.625, 0.75, 0.875,
# CHECK: 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 3.5,
# CHECK: 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0,
# CHECK: 16.0, 20.0, 24.0, 28.0, 32.0, 40.0, 48.0, 56.0,
# CHECK: 64.0, 80.0, 96.0, 112.0, 128.0, 160.0, 192.0, 224.0,
# CHECK: 256.0, 320.0, 384.0, 448.0, 512.0, 640.0, 768.0, 896.0,
# CHECK: 1024.0, 1280.0, 1536.0, 1792.0, 2048.0, 2560.0, 3072.0, 3584.0,
# CHECK: 4096.0, 5120.0, 6144.0, 7168.0, 8192.0, 10240.0, 12288.0, 14336.0,
# CHECK: 16384.0, 20480.0, 24576.0, 28672.0, 32768.0, 40960.0, 49152.0, 57344.0,
# CHECK: nan, -7.6293945e-06, -1.5258789e-05, -2.2888184e-05, -3.0517578e-05, -3.8146973e-05, -4.5776367e-05, -5.340576e-05,
# CHECK: -6.1035156e-05, -7.6293945e-05, -9.1552734e-05, -0.00010681152, -0.00012207031, -0.00015258789, -0.00018310547, -0.00021362305,
# CHECK: -0.00024414063, -0.00030517578, -0.00036621094, -0.0004272461, -0.00048828125, -0.00061035156, -0.0007324219, -0.0008544922,
# CHECK: -0.0009765625, -0.0012207031, -0.0014648438, -0.0017089844, -0.001953125, -0.0024414062, -0.0029296875, -0.0034179688,
# CHECK: -0.00390625, -0.0048828125, -0.005859375, -0.0068359375, -0.0078125, -0.009765625, -0.01171875, -0.013671875,
# CHECK: -0.015625, -0.01953125, -0.0234375, -0.02734375, -0.03125, -0.0390625, -0.046875, -0.0546875,
# CHECK: -0.0625, -0.078125, -0.09375, -0.109375, -0.125, -0.15625, -0.1875, -0.21875,
# CHECK: -0.25, -0.3125, -0.375, -0.4375, -0.5, -0.625, -0.75, -0.875,
# CHECK: -1.0, -1.25, -1.5, -1.75, -2.0, -2.5, -3.0, -3.5,
# CHECK: -4.0, -5.0, -6.0, -7.0, -8.0, -10.0, -12.0, -14.0,
# CHECK: -16.0, -20.0, -24.0, -28.0, -32.0, -40.0, -48.0, -56.0,
# CHECK: -64.0, -80.0, -96.0, -112.0, -128.0, -160.0, -192.0, -224.0,
# CHECK: -256.0, -320.0, -384.0, -448.0, -512.0, -640.0, -768.0, -896.0,
# CHECK: -1024.0, -1280.0, -1536.0, -1792.0, -2048.0, -2560.0, -3072.0, -3584.0,
# CHECK: -4096.0, -5120.0, -6144.0, -7168.0, -8192.0, -10240.0, -12288.0, -14336.0,
# CHECK: -16384.0, -20480.0, -24576.0, -28672.0, -32768.0, -40960.0, -49152.0, -57344.0,
fn test_simd_e5m2fnuz_to_f16():
    print("== test_simd_e5m2fnuz_to_f16")
    test_simd_e5m2fnuz_to_float[DType.float16]()


# CHECK-LABEL: test_simd_e5m2fnuz_to_bf16
# CHECK: 0.0, 7.6293945e-06, 1.5258789e-05, 2.2888184e-05, 3.0517578e-05, 3.8146973e-05, 4.5776367e-05, 5.340576e-05,
# CHECK: 6.1035156e-05, 7.6293945e-05, 9.1552734e-05, 0.00010681152, 0.00012207031, 0.00015258789, 0.00018310547, 0.00021362305,
# CHECK: 0.00024414063, 0.00030517578, 0.00036621094, 0.0004272461, 0.00048828125, 0.00061035156, 0.0007324219, 0.0008544922,
# CHECK: 0.0009765625, 0.0012207031, 0.0014648438, 0.0017089844, 0.001953125, 0.0024414062, 0.0029296875, 0.0034179688,
# CHECK: 0.00390625, 0.0048828125, 0.005859375, 0.0068359375, 0.0078125, 0.009765625, 0.01171875, 0.013671875,
# CHECK: 0.015625, 0.01953125, 0.0234375, 0.02734375, 0.03125, 0.0390625, 0.046875, 0.0546875,
# CHECK: 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.15625, 0.1875, 0.21875,
# CHECK: 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.625, 0.75, 0.875,
# CHECK: 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 3.5,
# CHECK: 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0,
# CHECK: 16.0, 20.0, 24.0, 28.0, 32.0, 40.0, 48.0, 56.0,
# CHECK: 64.0, 80.0, 96.0, 112.0, 128.0, 160.0, 192.0, 224.0,
# CHECK: 256.0, 320.0, 384.0, 448.0, 512.0, 640.0, 768.0, 896.0,
# CHECK: 1024.0, 1280.0, 1536.0, 1792.0, 2048.0, 2560.0, 3072.0, 3584.0,
# CHECK: 4096.0, 5120.0, 6144.0, 7168.0, 8192.0, 10240.0, 12288.0, 14336.0,
# CHECK: 16384.0, 20480.0, 24576.0, 28672.0, 32768.0, 40960.0, 49152.0, 57344.0,
# CHECK: nan, -7.6293945e-06, -1.5258789e-05, -2.2888184e-05, -3.0517578e-05, -3.8146973e-05, -4.5776367e-05, -5.340576e-05,
# CHECK: -6.1035156e-05, -7.6293945e-05, -9.1552734e-05, -0.00010681152, -0.00012207031, -0.00015258789, -0.00018310547, -0.00021362305,
# CHECK: -0.00024414063, -0.00030517578, -0.00036621094, -0.0004272461, -0.00048828125, -0.00061035156, -0.0007324219, -0.0008544922,
# CHECK: -0.0009765625, -0.0012207031, -0.0014648438, -0.0017089844, -0.001953125, -0.0024414062, -0.0029296875, -0.0034179688,
# CHECK: -0.00390625, -0.0048828125, -0.005859375, -0.0068359375, -0.0078125, -0.009765625, -0.01171875, -0.013671875,
# CHECK: -0.015625, -0.01953125, -0.0234375, -0.02734375, -0.03125, -0.0390625, -0.046875, -0.0546875,
# CHECK: -0.0625, -0.078125, -0.09375, -0.109375, -0.125, -0.15625, -0.1875, -0.21875,
# CHECK: -0.25, -0.3125, -0.375, -0.4375, -0.5, -0.625, -0.75, -0.875,
# CHECK: -1.0, -1.25, -1.5, -1.75, -2.0, -2.5, -3.0, -3.5,
# CHECK: -4.0, -5.0, -6.0, -7.0, -8.0, -10.0, -12.0, -14.0,
# CHECK: -16.0, -20.0, -24.0, -28.0, -32.0, -40.0, -48.0, -56.0,
# CHECK: -64.0, -80.0, -96.0, -112.0, -128.0, -160.0, -192.0, -224.0,
# CHECK: -256.0, -320.0, -384.0, -448.0, -512.0, -640.0, -768.0, -896.0,
# CHECK: -1024.0, -1280.0, -1536.0, -1792.0, -2048.0, -2560.0, -3072.0, -3584.0,
# CHECK: -4096.0, -5120.0, -6144.0, -7168.0, -8192.0, -10240.0, -12288.0, -14336.0,
# CHECK: -16384.0, -20480.0, -24576.0, -28672.0, -32768.0, -40960.0, -49152.0, -57344.0,
fn test_simd_e5m2fnuz_to_bf16():
    print("== test_simd_e5m2fnuz_to_bf16")
    test_simd_e5m2fnuz_to_float[DType.bfloat16]()


# CHECK-LABEL: test_simd_f32_to_e5m2fnuz
# CHECK: -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0,
# CHECK: -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0,
# CHECK: -256.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0,
# CHECK: -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0,
# CHECK: -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0,
# CHECK: -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0,
# CHECK: -96.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0,
# CHECK: -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0,
# CHECK: -64.0, -64.0, -64.0, -64.0, -64.0, -64.0, -64.0, -64.0,
# CHECK: -64.0, -64.0, -64.0, -64.0, -64.0, -56.0, -56.0, -56.0,
# CHECK: -56.0, -56.0, -56.0, -56.0, -48.0, -48.0, -48.0, -48.0,
# CHECK: -48.0, -48.0, -48.0, -48.0, -48.0, -40.0, -40.0, -40.0,
# CHECK: -40.0, -40.0, -40.0, -40.0, -32.0, -32.0, -32.0, -32.0,
# CHECK: -32.0, -32.0, -32.0, -28.0, -28.0, -28.0, -24.0, -24.0,
# CHECK: -24.0, -24.0, -24.0, -20.0, -20.0, -20.0, -16.0, -16.0,
# CHECK: -16.0, -16.0, -14.0, -12.0, -12.0, -12.0, -10.0, -8.0,
# CHECK: -8.0, -7.0, -6.0, -5.0, -4.0, -3.0, -2.0, -1.0,
# CHECK: 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
# CHECK: 8.0, 8.0, 10.0, 12.0, 12.0, 12.0, 14.0, 16.0,
# CHECK: 16.0, 16.0, 16.0, 20.0, 20.0, 20.0, 24.0, 24.0,
# CHECK: 24.0, 24.0, 24.0, 28.0, 28.0, 28.0, 32.0, 32.0,
# CHECK: 32.0, 32.0, 32.0, 32.0, 32.0, 40.0, 40.0, 40.0,
# CHECK: 40.0, 40.0, 40.0, 40.0, 48.0, 48.0, 48.0, 48.0,
# CHECK: 48.0, 48.0, 48.0, 48.0, 48.0, 56.0, 56.0, 56.0,
# CHECK: 56.0, 56.0, 56.0, 56.0, 64.0, 64.0, 64.0, 64.0,
# CHECK: 64.0, 64.0, 64.0, 64.0, 64.0, 64.0, 64.0, 64.0,
# CHECK: 64.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0,
# CHECK: 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0,
# CHECK: 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0,
# CHECK: 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0,
# CHECK: 96.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0,
# CHECK: 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0,
# CHECK: 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0,
fn test_simd_f32_to_e5m2fnuz():
    print("== test_simd_f32_to_e5m2fnuz")

    comptime M = 512
    var f32_simd = SIMD[DType.float32, M](0.0)

    for i in range(M):
        f32_simd[i] = i - 256

    f32_casted_e5m2 = f32_simd.cast[DType.float8_e5m2fnuz]()

    for i in range(64):
        for j in range(8):
            print(f32_casted_e5m2[i * 8 + j], end=", ")
        print("")


# CHECK-LABEL: test_simd_e5m2fnuz_to_f32_ptx_path
# CHECK: 0.0, 7.6293945e-06, 1.5258789e-05, 2.2888184e-05, 3.0517578e-05, 3.8146973e-05, 4.5776367e-05, 5.340576e-05,
# CHECK: 6.1035156e-05, 7.6293945e-05, 9.1552734e-05, 0.00010681152, 0.00012207031, 0.00015258789, 0.00018310547, 0.00021362305,
# CHECK: 0.00024414063, 0.00030517578, 0.00036621094, 0.0004272461, 0.00048828125, 0.00061035156, 0.0007324219, 0.0008544922,
# CHECK: 0.0009765625, 0.0012207031, 0.0014648438, 0.0017089844, 0.001953125, 0.0024414062, 0.0029296875, 0.0034179688,
# CHECK: 0.00390625, 0.0048828125, 0.005859375, 0.0068359375, 0.0078125, 0.009765625, 0.01171875, 0.013671875,
# CHECK: 0.015625, 0.01953125, 0.0234375, 0.02734375, 0.03125, 0.0390625, 0.046875, 0.0546875,
# CHECK: 0.0625, 0.078125, 0.09375, 0.109375, 0.125, 0.15625, 0.1875, 0.21875,
# CHECK: 0.25, 0.3125, 0.375, 0.4375, 0.5, 0.625, 0.75, 0.875,
# CHECK: 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0, 3.5,
# CHECK: 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0,
# CHECK: 16.0, 20.0, 24.0, 28.0, 32.0, 40.0, 48.0, 56.0,
# CHECK: 64.0, 80.0, 96.0, 112.0, 128.0, 160.0, 192.0, 224.0,
# CHECK: 256.0, 320.0, 384.0, 448.0, 512.0, 640.0, 768.0, 896.0,
# CHECK: 1024.0, 1280.0, 1536.0, 1792.0, 2048.0, 2560.0, 3072.0, 3584.0,
# CHECK: 4096.0, 5120.0, 6144.0, 7168.0, 8192.0, 10240.0, 12288.0, 14336.0,
# CHECK: 16384.0, 20480.0, 24576.0, 28672.0, 32768.0, 40960.0, 49152.0, 57344.0,
# CHECK: nan, -7.6293945e-06, -1.5258789e-05, -2.2888184e-05, -3.0517578e-05, -3.8146973e-05, -4.5776367e-05, -5.340576e-05,
# CHECK: -6.1035156e-05, -7.6293945e-05, -9.1552734e-05, -0.00010681152, -0.00012207031, -0.00015258789, -0.00018310547, -0.00021362305,
# CHECK: -0.00024414063, -0.00030517578, -0.00036621094, -0.0004272461, -0.00048828125, -0.00061035156, -0.0007324219, -0.0008544922,
# CHECK: -0.0009765625, -0.0012207031, -0.0014648438, -0.0017089844, -0.001953125, -0.0024414062, -0.0029296875, -0.0034179688,
# CHECK: -0.00390625, -0.0048828125, -0.005859375, -0.0068359375, -0.0078125, -0.009765625, -0.01171875, -0.013671875,
# CHECK: -0.015625, -0.01953125, -0.0234375, -0.02734375, -0.03125, -0.0390625, -0.046875, -0.0546875,
# CHECK: -0.0625, -0.078125, -0.09375, -0.109375, -0.125, -0.15625, -0.1875, -0.21875,
# CHECK: -0.25, -0.3125, -0.375, -0.4375, -0.5, -0.625, -0.75, -0.875,
# CHECK: -1.0, -1.25, -1.5, -1.75, -2.0, -2.5, -3.0, -3.5,
# CHECK: -4.0, -5.0, -6.0, -7.0, -8.0, -10.0, -12.0, -14.0,
# CHECK: -16.0, -20.0, -24.0, -28.0, -32.0, -40.0, -48.0, -56.0,
# CHECK: -64.0, -80.0, -96.0, -112.0, -128.0, -160.0, -192.0, -224.0,
# CHECK: -256.0, -320.0, -384.0, -448.0, -512.0, -640.0, -768.0, -896.0,
# CHECK: -1024.0, -1280.0, -1536.0, -1792.0, -2048.0, -2560.0, -3072.0, -3584.0,
# CHECK: -4096.0, -5120.0, -6144.0, -7168.0, -8192.0, -10240.0, -12288.0, -14336.0,
# CHECK: -16384.0, -20480.0, -24576.0, -28672.0, -32768.0, -40960.0, -49152.0, -57344.0,
fn test_simd_e5m2fnuz_to_f32_ptx_path(ctx: DeviceContext) raises:
    print("== test_simd_e5m2fnuz_to_f32_ptx_path")

    comptime M = 256
    var e5m2_simd = SIMD[DType.float8_e5m2fnuz, M](0.0)
    for i in range(M):
        e5m2_simd[i] = bitcast[DType.float8_e5m2fnuz](UInt8(i))

    comptime kernel = test_simd_float8[DType.float8_e5m2fnuz, M, DType.float32]
    ctx.enqueue_function_experimental[kernel](
        e5m2_simd, grid_dim=1, block_dim=1
    )
    ctx.synchronize()


fn test_simd_float32[
    size: Int,
    target: DType,
](x: SIMD[DType.float32, size]):
    var x_casted = x.cast[target]()

    comptime M = 64
    comptime N = size // M
    for i in range(M):
        for j in range(N):
            print(x_casted[i * N + j], end=", ")
        print("")


# CHECK-LABEL: test_simd_f32_to_e5m2fnuz_ptx_path
# CHECK: -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0,
# CHECK: -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0, -256.0,
# CHECK: -256.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0, -224.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0, -192.0,
# CHECK: -192.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0, -160.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0, -128.0,
# CHECK: -128.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0,
# CHECK: -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0, -112.0,
# CHECK: -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0,
# CHECK: -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0, -96.0,
# CHECK: -96.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0,
# CHECK: -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0, -80.0,
# CHECK: -64.0, -64.0, -64.0, -64.0, -64.0, -64.0, -64.0, -64.0,
# CHECK: -64.0, -64.0, -64.0, -64.0, -64.0, -56.0, -56.0, -56.0,
# CHECK: -56.0, -56.0, -56.0, -56.0, -48.0, -48.0, -48.0, -48.0,
# CHECK: -48.0, -48.0, -48.0, -48.0, -48.0, -40.0, -40.0, -40.0,
# CHECK: -40.0, -40.0, -40.0, -40.0, -32.0, -32.0, -32.0, -32.0,
# CHECK: -32.0, -32.0, -32.0, -28.0, -28.0, -28.0, -24.0, -24.0,
# CHECK: -24.0, -24.0, -24.0, -20.0, -20.0, -20.0, -16.0, -16.0,
# CHECK: -16.0, -16.0, -14.0, -12.0, -12.0, -12.0, -10.0, -8.0,
# CHECK: -8.0, -7.0, -6.0, -5.0, -4.0, -3.0, -2.0, -1.0,
# CHECK: 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0,
# CHECK: 8.0, 8.0, 10.0, 12.0, 12.0, 12.0, 14.0, 16.0,
# CHECK: 16.0, 16.0, 16.0, 20.0, 20.0, 20.0, 24.0, 24.0,
# CHECK: 24.0, 24.0, 24.0, 28.0, 28.0, 28.0, 32.0, 32.0,
# CHECK: 32.0, 32.0, 32.0, 32.0, 32.0, 40.0, 40.0, 40.0,
# CHECK: 40.0, 40.0, 40.0, 40.0, 48.0, 48.0, 48.0, 48.0,
# CHECK: 48.0, 48.0, 48.0, 48.0, 48.0, 56.0, 56.0, 56.0,
# CHECK: 56.0, 56.0, 56.0, 56.0, 64.0, 64.0, 64.0, 64.0,
# CHECK: 64.0, 64.0, 64.0, 64.0, 64.0, 64.0, 64.0, 64.0,
# CHECK: 64.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0,
# CHECK: 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0, 80.0,
# CHECK: 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0,
# CHECK: 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0, 96.0,
# CHECK: 96.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0,
# CHECK: 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0, 112.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0, 128.0,
# CHECK: 128.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0, 160.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0, 192.0,
# CHECK: 192.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0, 224.0,
# CHECK: 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0,
# CHECK: 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0, 256.0,
fn test_simd_f32_to_e5m2fnuz_ptx_path(ctx: DeviceContext) raises:
    print("== test_simd_f32_to_e5m2fnuz_ptx_path")

    comptime M = 512
    var f32_simd = SIMD[DType.float32, M](0.0)
    for i in range(M):
        f32_simd[i] = i - 256

    comptime kernel = test_simd_float32[M, DType.float8_e5m2fnuz]
    ctx.enqueue_function_experimental[kernel](f32_simd, grid_dim=1, block_dim=1)
    ctx.synchronize()


fn test_simd_float8[
    dtype: DType,
    size: Int,
    target: DType,
](x: SIMD[dtype, size]):
    var x_casted = x.cast[target]()

    comptime M = 32
    comptime N = size // M
    for i in range(M):
        for j in range(N):
            print(x_casted[i * N + j], end=", ")
        print("")


def main():
    test_e5m2fnuz_initialization()

    test_simd_e5m2fnuz_to_f32()
    test_simd_e5m2fnuz_to_f16()
    test_simd_e5m2fnuz_to_bf16()

    test_simd_f32_to_e5m2fnuz()

    with DeviceContext() as ctx:
        test_simd_e5m2fnuz_to_f32_ptx_path(ctx)
        test_simd_f32_to_e5m2fnuz_ptx_path(ctx)
